#include <cstring>
#include "lit/elf/elfFile.h"
#include "nasm/type/type.h"
#include "lit/elf/symLink.h"

void ElfFile::readElf(const std::string &fileName) {
    this->fileName = fileName;
    // 打开目标文件
    FILE *file = fopen(fileName.c_str(), "rb");

    // 重置文件读取位置为文件头
    rewind(file);
    fread(&ehdr, sizeof(Elf32_Ehdr), 1, file);

    // 如果是可执行文件，读取程序头表
    if (ehdr.e_type == ET_EXEC) {
        fseek(file, ehdr.e_phoff, 0);
        for (u32 i = 0; i < ehdr.e_phnum; ++i) {
            std::shared_ptr<Elf32_Phdr> phdr = std::make_shared<Elf32_Phdr>();
            fread(phdr.get(), ehdr.e_phentsize, 1, file);
            phdrTable.emplace_back(phdr);
        }
    }

    // .shstrtab头部表
    Elf32_Shdr sh_shstrTable;
    fseek(file, ehdr.e_shoff + ehdr.e_shentsize * ehdr.e_shstrndx, 0);
    fread(&sh_shstrTable, ehdr.e_shentsize, 1, file);

    // .shstrtab
    char shstrtabData[sh_shstrTable.sh_size];
    fseek(file, sh_shstrTable.sh_offset, 0);
    fread(shstrtabData, sh_shstrTable.sh_size, 1, file);
    shstrtab = shstrtabData;

    // 段表
    fseek(file, ehdr.e_shoff, 0);
    for (u32 i = 0; i < ehdr.e_shnum; ++i) {
        std::shared_ptr<Elf32_Shdr> shdr = std::make_shared<Elf32_Shdr>();
        fread(shdr.get(), ehdr.e_shentsize, 1, file);
        std::string name = shstrtabData + shdr->sh_name;
        shdrNames.emplace_back(name);
        shdrTable[name] = shdr;
    }

    // .strtab
    std::shared_ptr<Elf32_Shdr> sh_strtab = shdrTable[".strtab"];
    char strTabData[sh_strtab->sh_size];
    fseek(file, sh_strtab->sh_offset, 0);
    fread(strTabData, sh_strtab->sh_size, 1, file);
    strtab = strTabData;

    // .symtab
    std::shared_ptr<Elf32_Shdr> sh_symTab = shdrTable[".symtab"];
    fseek(file, sh_symTab->sh_offset, 0);
    u32 symNum = sh_symTab->sh_size / sh_symTab->sh_entsize;
    for (u32 i = 0; i < symNum; ++i) {
        std::shared_ptr<Elf32_Sym> sym = std::make_shared<Elf32_Sym>();
        fread(sym.get(), sh_symTab->sh_entsize, 1, file);
        std::string name = strTabData + sym->st_name;
        symbolNames.emplace_back(name);
        symbolTable[name] = sym;
    }

    // .rel.data .rel.text
    for (auto &shdrName: shdrNames) {
        auto shdr = shdrTable[shdrName];
        if (shdr->sh_type == SHT_REL) {
            fseek(file, shdr->sh_offset, 0);
            u32 relNum = shdr->sh_size / shdr->sh_entsize;
            for (u32 j = 0; j < relNum; ++j) {
                auto rel = std::make_shared<Elf32_Rel>();
                fread(rel.get(), shdr->sh_entsize, 1, file);
                std::string segName = shdrNames[shdr->sh_info];
                std::string symName = symbolNames[ELF32_R_SYM(rel->r_info)];
                relTable.emplace_back(std::make_shared<RelocationItem>(segName, rel, symName));
            }
        }
    }
    fclose(file);
}

void ElfFile::getData(char *buffer, Elf32_Off offset, Elf32_Word size) const {
    FILE *file = fopen(fileName.c_str(), "rb");
    rewind(file);
    fseek(file, offset, 0);
    fread(buffer, size, 1, file);
    fclose(file);
}

void ElfFile::assemObj(Linker *linker) {
    // 记录生成elf文件的所有段名
    std::vector<std::string> allSegNames;
    allSegNames.emplace_back("");
    std::vector<std::string> segNames = linker->segNames;
    for (auto &segName: segNames) {
        allSegNames.emplace_back(segName);
    }
    allSegNames.emplace_back(".shstrtab");
    allSegNames.emplace_back(".symtab");
    allSegNames.emplace_back(".strtab");

    // 段索引
    std::unordered_map<std::string, u32> shIndex;
    // 段名索引
    std::unordered_map<std::string, u32> shstrIndex;
    // 遍历所有的段，更新段索引以及段名索引
    for (u32 i = 0; i < allSegNames.size(); ++i) {
        std::string name = allSegNames[i];
        shIndex[name] = i;
        shstrIndex[name] = shstrIndex.size();
        shstrtab += name;
        shstrtab.push_back('\0');
    }

    // 添加一个新符号到符号表中
    addSym("", nullptr);
    // 获取所有定义的符号
    auto symDef = linker->symDef;
    // 遍历所有定义的符号
    for (auto &def: symDef) {
        std::string name = def->name;       // 获取当前遍历的符号名
        std::shared_ptr<ElfFile> prov = def->prov;  // 获取符号的提供者
        auto sym = prov->symbolTable[name]; // 获取当前遍历的符号
        std::string segName = prov->shdrNames[sym->st_shndx]; // 获取符号所在的段名
        sym->st_shndx = shIndex[segName];  // 设置符号的段索引
        addSym(name, sym);  // 添加当前遍历的符号对象到可执行文件的符号表中
    }

    // 符号索引
    std::unordered_map<std::string, u32> symIndex;
    // 符号名索引
    std::unordered_map<std::string, u32> strIndex;
    // 更新符号名索引,符号索引
    for (u32 i = 0; i < symbolNames.size(); ++i) {
        std::string name = symbolNames[i];  // 获取符号名
        symIndex[name] = i;     // 记录符号索引
        strIndex[name] = strtab.size(); // 记录符号名索引
        strtab += name; // 添加符号名到字符串表中
        strtab.push_back('\0');
    }
    // 更新可执行文件中的符号表中的符号名索引
    for (auto &name: symbolNames) {
        symbolTable[name]->st_name = strIndex[name];
    }

    // elf文件头
    char magic[] = {
            0x7f, 0x45, 0x4c, 0x46, 0x01, 0x01, 0x01, 0x00,
            0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00
    };

    memcpy(&ehdr.e_ident, magic, sizeof(magic));
    ehdr.e_type = ET_EXEC;
    ehdr.e_machine = EM_386;
    ehdr.e_version = EV_CURRENT;
    ehdr.e_entry = symbolTable[START]->st_value;
    ehdr.e_phoff = 0;
    ehdr.e_shoff = 0;
    ehdr.e_flags = 0;
    ehdr.e_ehsize = sizeof(Elf32_Ehdr);
    ehdr.e_phentsize = sizeof(Elf32_Phdr);
    ehdr.e_phnum = segNames.size();
    ehdr.e_shentsize = sizeof(Elf32_Shdr);
    ehdr.e_shnum = allSegNames.size();
    ehdr.e_shstrndx = shIndex[".shstrtab"];

    u32 curOff = sizeof(ehdr);
    // 设置程序头部表的偏移
    ehdr.e_phoff = curOff;
    // 生成程序头部表
    for (auto &name: segNames) {
        Elf32_Word flags = PF_W | PF_R;
        if (name == ".text") {
            flags = PF_X | PF_R;
        }
        addPhdr(PT_LOAD, linker->segLists[name]->offset, linker->segLists[name]->baseAddr,
                linker->segLists[name]->size, linker->segLists[name]->size, flags, MEM_ALIGN);
    }

    // 更新curOff为段表偏移
    curOff += ehdr.e_phentsize * ehdr.e_phnum;
    // 生成程序的段表
    addShdr("", 0, 0, 0, 0, 0, 0, 0, 0, 0);
    // 生成.text和.data段
    for (auto &name: segNames) {
        Elf32_Word sh_flags = SHF_ALLOC | SHF_WRITE;
        Elf32_Word sh_align = DISC_ALIGN;
        if (name == ".text") {
            sh_flags = SHF_ALLOC | SHF_EXECINSTR;
            sh_align = TEXT_ALIGN;
        }
        addShdr(name, SHT_PROGBITS, sh_flags, linker->segLists[name]->baseAddr, linker->segLists[name]->offset,
                linker->segLists[name]->size, 0, 0, sh_align, 0);
        curOff = linker->segLists[name]->offset + linker->segLists[name]->size;
    }
    curOff += (4 - curOff % 4) % 4;
    // 添加段名字符串表
    addShdr(".shstrtab", SHT_STRTAB, 0, 0, curOff, shstrtab.size(), SHN_UNDEF, 0, 1, 0);
    curOff += shstrtab.size();
    curOff += (4 - curOff % 4) % 4;
    // curOff增加段表偏移
    ehdr.e_shoff = curOff;
    curOff += ehdr.e_shnum * ehdr.e_shentsize;
    // 添加符号表
    addShdr(".symtab", SHT_SYMTAB, 0, 0, curOff, symbolNames.size() * sizeof(Elf32_Sym), shIndex[".strtab"], 0, 1,
            sizeof(Elf32_Sym));
    curOff += symbolNames.size() * sizeof(Elf32_Sym);
    // 添加字符串表
    addShdr(".strtab", SHT_STRTAB, 0, 0, curOff, strtab.size(), SHN_UNDEF, 0, 1, 0);
    curOff += strtab.size();
    curOff += (4 - curOff % 4) % 4;

    // 更新段表段名索引
    for (auto &name: allSegNames) {
        shdrTable[name]->sh_name = shstrIndex[name];
    }

}

void ElfFile::addSym(const std::string &symName, std::shared_ptr<Elf32_Sym> sym) {
    std::shared_ptr<Elf32_Sym> newSym = std::make_shared<Elf32_Sym>();
    if (symName.empty()) {
        newSym->st_name = 0;
        newSym->st_value = 0;
        newSym->st_size = 0;
        newSym->st_info = 0;
        newSym->st_other = 0;
        newSym->st_shndx = 0;
    } else {
        newSym->st_name = 0;
        newSym->st_value = sym->st_value;
        newSym->st_size = sym->st_size;
        newSym->st_info = sym->st_info;
        newSym->st_other = sym->st_other;
        newSym->st_shndx = sym->st_shndx;
    }
    symbolTable[symName] = newSym;
    symbolNames.emplace_back(symName);
}

void ElfFile::addPhdr(Elf32_Word type, Elf32_Off off, Elf32_Addr vaddr, Elf32_Word filesz, Elf32_Word memsz,
                      Elf32_Word flags, Elf32_Word align) {
    auto ph = std::make_shared<Elf32_Phdr>();
    ph->p_type = type;
    ph->p_offset = off;
    ph->p_vaddr = ph->p_paddr = vaddr;
    ph->p_filesz = filesz;
    ph->p_memsz = memsz;
    ph->p_flags = flags;
    ph->p_align = align;
    phdrTable.push_back(ph);
}

void
ElfFile::addShdr(std::string sh_name, Elf32_Word sh_type, Elf32_Word sh_flags, Elf32_Addr sh_addr, Elf32_Off sh_offset,
                 Elf32_Word sh_size, Elf32_Word sh_link, Elf32_Word sh_info, Elf32_Word sh_addralign,
                 Elf32_Word sh_entsize) {
    auto sh = std::make_shared<Elf32_Shdr>();
    sh->sh_name = 0;
    sh->sh_type = sh_type;
    sh->sh_flags = sh_flags;
    sh->sh_addr = sh_addr;
    sh->sh_offset = sh_offset;
    sh->sh_size = sh_size;
    sh->sh_link = sh_link;
    sh->sh_info = sh_info;
    sh->sh_addralign = sh_addralign;
    sh->sh_entsize = sh_entsize;
    shdrTable[sh_name] = sh;
    shdrNames.push_back(sh_name);
}

void ElfFile::writeElf(Linker *linker, const std::string &fileName) {
    FILE *fileOut = fopen(fileName.c_str(), "wb");
    u32 padNum = 0;
    char pad[1] = {0};
    // 输出elf文件头
    fwrite(&ehdr, ehdr.e_ehsize, 1, fileOut);
    // 输出elf文件头部表
    for (auto &phdr: phdrTable) {
        fwrite(phdr.get(), ehdr.e_phentsize, 1, fileOut);
    }
    // 输出.text和.data段
    for (u32 i = 0; i < linker->segNames.size(); ++i) {
        auto segs = linker->segLists[linker->segNames[i]];
        padNum = segs->offset - segs->begin;
        // 输出填充字段
        fwrite(pad, sizeof(pad), padNum, fileOut);
        // 记录上一个数据块
        std::shared_ptr<Block> last;
        // 遍历当前段的所有数据块
        for (u32 j = 0; j < segs->blocks.size(); ++j) {
            auto block = segs->blocks[j];
            if (last != nullptr) {
                // 如果当前输出的数据块不是第一个数据块
                u32 lastEnd = last->offset + last->size;    // 计算上一个数据块输出的结束位置
                padNum = block->offset - lastEnd;  // 计算输出当前数据块之前需要填充的0字节数量
                // 输出填充字段
                fwrite(block->data, sizeof(pad), padNum, fileOut);
            }
            // 输出当前数据块
            fwrite(block->data, block->size, 1, fileOut);
            last = block;
        }
    }
    // 计算.shstrtab与.data段之间需要填充的字节数量
    padNum = shdrTable[".shstrtab"]->sh_offset - shdrTable[".data"]->sh_offset - shdrTable[".data"]->sh_size;
    // 输出填充字段
    fwrite(pad, sizeof(pad), padNum, fileOut);
    // 输出.shstrtab段
    fwrite(shstrtab.c_str(), shstrtab.size(), 1, fileOut);
    // 计算section header table与shstrtab之间需要填充的字节数量
    padNum = ehdr.e_shoff - shdrTable[".shstrtab"]->sh_offset - shdrTable[".shstrtab"]->sh_size;
    // 输出填充字段
    fwrite(pad, sizeof(pad), padNum, fileOut);
    // 输出section header table
    for (const auto &shdrName: shdrNames) {
        auto shdr = shdrTable[shdrName];
        fwrite(shdr.get(), ehdr.e_shentsize, 1, fileOut);
    }

    // 输出符号表
    for (const auto &symbolName: symbolNames) {
        auto sym = symbolTable[symbolName];
        fwrite(sym.get(), sizeof(Elf32_Sym), 1, fileOut);
    }
    // 输出字符串表
    fwrite(shstrtab.c_str(), shstrtab.size(), 1, fileOut);
    fclose(fileOut);
}
